﻿#include "MaterialGraphModel.h"

#include "MaterialGraph/Link/MaterialGraphLinkModel.h"
#include "MaterialGraph/Node/MaterialGraphNodeModel.h"
#include "MaterialGraph/Pin/MaterialGraphPinModel.h"

#include "Object/NewObject.h"

#include "MaterialGraph/MaterialGraphModel.h"
#include "MaterialGraph/BRDFNode/MaterialBRDFNodeModel.h"

#include "Math/Math.h"
#include "Output/Console.h"
#include <stack>

#include "Framework/Application.h"
#include "MaterialGraph/ShaderCompileSubsystem.h"

#include "Material/Material.h"

#include "Asset/AssetSubsystem.h"

bool RMaterialGraphModel::CanUndo() const
{
    return LastOperationIndex>=0;
}

bool RMaterialGraphModel::CanRedo() const
{
    return LastOperationIndex<(int64_t)(Operations.Num())-1;
}

void RMaterialGraphModel::Undo()
{
    if(LastOperationIndex<0)
    {
        return;
    }

    Operations[LastOperationIndex]->Undo(this);
    LastOperationIndex--;
}

void RMaterialGraphModel::Redo()
{
    if(Operations.Empty())
    {
        return;
    }

    if(LastOperationIndex<(int64_t)(Operations.Num())-1)
    {
        Operations[LastOperationIndex+1]->Redo(this);
        LastOperationIndex++;
    }
}

int64_t RMaterialGraphModel::GetNodeNum() const
{
    return Nodes.Num();
}

TSharedObjectPtr<RMaterialGraphNodeModel> RMaterialGraphModel::GetNode(int64_t Index) const
{
    return Nodes[Index];
}

void RMaterialGraphModel::AddNode(TSubclassOf<RMaterialGraphNodeModel> NodeClass, const SVec2 &NodePosition)
{
    auto Node=NewObject<RMaterialGraphNodeModel>(NodeClass.Get(),this);
    Node->Position=NodePosition;

    AddNode(Node);
}

void RMaterialGraphModel::AddNode(TSharedObjectPtr<RMaterialGraphNodeModel> Node)
{
    auto AddNodeOperation=NewObject<RMaterialAddNodeOperation>();
    AddNodeOperation->Node=Node;
    AddNodeOperation->Position=Node->Position;
    AddNodeOperation->NodeIndex=Nodes.Num();

    AddOperation(AddNodeOperation);
    AddNodeOperation->Redo(this);

    //如果是输出节点，那么将其设置为当前输出节点
    if(Node->IsOutputNode())
    {
        SetCurrentOutputNode(Node);
    }

}

void RMaterialGraphModel::RemoveNode(TSharedObjectPtr<RMaterialGraphNodeModel> Node)
{
    auto Index=Nodes.Find(Node).GetValue();
    auto RemoveNodeOperation=NewObject<RMaterialRemoveNodeOperation>();
    RemoveNodeOperation->Node=Node;
    RemoveNodeOperation->NodeIndex=Index;

    AddOperation(RemoveNodeOperation);
    RemoveNodeOperation->Redo(this);

    if(SelectedNode==Node)
    {
        SetSelectedNode(nullptr);
    }

    //如果删除的是当前输出节点，那么将其设置为nullptr
    if(CurrentOutputNode==Node)
    {
        SetCurrentOutputNode(nullptr);
    }
}

int64_t RMaterialGraphModel::GetLinkNum() const
{
    return Links.Num();
}

TSharedObjectPtr<RMaterialGraphLinkModel> RMaterialGraphModel::GetLink(int64_t Index) const
{
    return Links[Index];
}

void RMaterialGraphModel::AddLink(TSharedObjectPtr<RMaterialGraphLinkModel> Link)
{
    auto AddLinkOperation=NewObject<RMaterialAddLinkOperation>();
    AddLinkOperation->AddedLink=Link;
    AddLinkOperation->LinkIndex=Links.Num();

    AddOperation(AddLinkOperation);

    AddLinkOperation->Redo(this);

}

void RMaterialGraphModel::RemoveLink(TSharedObjectPtr<RMaterialGraphLinkModel> Link)
{
    auto Index=Links.Find(Link);

    auto RemoveLinkOperation=NewObject<RMaterialRemoveLinkOperation>();
    RemoveLinkOperation->RemovedLink=Link;
    RemoveLinkOperation->LinkIndex=Index.GetValue();

    AddOperation(RemoveLinkOperation);

    RemoveLinkOperation->Redo(this);
}

TSharedObjectPtr<RMaterialGraphLinkModel> RMaterialGraphModel::CreateAndAddLink(TSharedObjectPtr<RMaterialGraphNodeModel> StartNode, int64_t InOuputPinIndex, TSharedObjectPtr<RMaterialGraphNodeModel> EndNode, int64_t InInputPinIndex)
{
    auto Link=NewObject<RMaterialGraphLinkModel>(this);
    Link->Start=StartNode->GetOutput(InOuputPinIndex);
    Link->End=EndNode->GetInput(InInputPinIndex);

    // StartNode->GetOutput(InOuputPinIndex)->Links.Add(Link);
    // EndNode->GetInput(InInputPinIndex)->Links.Add(Link);

    AddLink(Link);

    return Link;
}

TSharedObjectPtr<RMaterialGraphLinkModel> RMaterialGraphModel::CreateAndAddLink(TSharedObjectPtr<RMaterialGraphPinModel> StartPin, TSharedObjectPtr<RMaterialGraphPinModel> EndPin)
{
    auto Link=NewObject<RMaterialGraphLinkModel>(this);
    Link->Start=StartPin;
    Link->End=EndPin;

    // assert(StartPin->bIsOutput);
    // assert(!EndPin->bIsOutput);

    // StartPin->Links.Add(Link);
    // EndPin->Links.Add(Link);

    AddLink(Link);

    return Link;
}


void RMaterialGraphModel::RegenerateMaterial(TSharedObjectPtr<RMaterial> InMaterial)
{
    //check output node in the graph
    if(!CurrentOutputNode)
    {
        WH::PrintLine(U"材质更新失败： 材质必须有一个输出节点");
        return;
    }

    //sort the nodes,output node must be the last node
    TVector<TSharedObjectPtr<RMaterialGraphNodeModel>> OrderedNodes;
    {
        //1. calculate the depth of the nodes
        TMap<TSharedObjectPtr<RMaterialGraphNodeModel>, int64_t> NodeDepthMap;
        NodeDepthMap.Add(CurrentOutputNode, 0);

        std::stack<TSharedObjectPtr<RMaterialGraphNodeModel>> NodeStacks;
        NodeStacks.push(CurrentOutputNode);

        while(!NodeStacks.empty())
        {
            auto Node =  NodeStacks.top();
            NodeStacks.pop();

            int64_t Depth = NodeDepthMap[Node];

            for(int64_t i=0; i<Node->GetInputNum(); ++i)
            {
                auto Pin = Node->GetInput(i);
                for(auto& Link : Pin->GetLinks())
                {
                    if(!Link->IsValid())
                    {
                        continue;
                    }

                    auto StartNode = Link->Start->GetNode();
                    if(NodeDepthMap.Contains(StartNode))
                    {
                        int64_t OldDepth = NodeDepthMap[StartNode];
                        NodeDepthMap.Add(StartNode, CMath::Max(OldDepth,Depth+1) );
                    }
                    else
                    {
                        NodeDepthMap.Add(StartNode, Depth+1);
                    }
                    NodeStacks.push(StartNode);
                }
            }
        }

        //2.sort the nodes by depth
        //used the linked nodes,rather than all nodes
        OrderedNodes= NodeDepthMap.Keys();
        OrderedNodes.Sort([&NodeDepthMap](const TSharedObjectPtr<RMaterialGraphNodeModel>& A, const TSharedObjectPtr<RMaterialGraphNodeModel>& B)
        {
            return NodeDepthMap[A] > NodeDepthMap[B];
        });
    }

    //give every linked output pin a unique name
    //TMap<TSharedObjectPtr<RMaterialGraphPinModel>, CString> OutputPinNameMap;
    {
        int64_t OutputPinIndex = 0;
        for(auto& Node : OrderedNodes)
        {
            for(int64_t i=0; i<Node->GetOutputNum(); ++i)
            {
                auto Pin = Node->GetOutput(i);
                if(!Pin->HasValidLink())
                {
                    continue;
                }

                Pin->UniqueOutputPinName=U"Output"+CString(OutputPinIndex);
                OutputPinIndex++;
                //OutputPinNameMap.Add(Pin, CString(U"Output")+CString(OutputPinIndex));
            }
        }
    }

    //generate the shader code
    //generate text  
    CString FileText= CurrentOutputNode->GetShaderTemplate();

    TVector<float> FloatParameters ;
    TVector<TSharedObjectPtr<RTexture2D>> Texture2Ds ;
    TMap<CString,SMaterialParameterDescription> ParameterDescriptions ;

    //收集参数和纹理
    {
        for(auto& Node: OrderedNodes)
        {
            Node->CollectParameterAndTextures(FloatParameters,Texture2Ds,ParameterDescriptions);
        }

        // if(Textures.Num()>0)
        // {
        //     IncludeHeaders+=U"layout (binding = 0) uniform sampler2D Texture2Ds[];\n";
        // }
    }

    //在着色器中引入参数
    CString IncludeHeaders; // include shader functions
    {
        if(Texture2Ds.Num()>0)
        {
            IncludeHeaders+=U"layout (binding = 0) uniform sampler2D Texture2Ds[];\n";
        }

        if(FloatParameters.Num()>0)
        {
            int BindingIndex=Texture2Ds.Num()>0 ? 1 :0;
            IncludeHeaders+=U"layout (binding = "+CString(BindingIndex)+U") uniform SFloatParameters {\n";
            IncludeHeaders+=U"    vec4 Floats[];\n";
            IncludeHeaders+=U"} FloatParameters;\n";
        }
    }

    CString Calculation;//invoke function and store the result 
    for(auto& MaterialNode: OrderedNodes)
    {
        IncludeHeaders+=MaterialNode->GetIncludeFiles();
        Calculation+=MaterialNode->GenerateExpression();
    }


    FileText=FileText.Replace(U"/*HEADERS_AREA*/",IncludeHeaders);
    FileText=FileText.Replace(U"/*CALCULATE_AREA*/",Calculation);    

    WH::PrintLine(U"Shader Code:");
    WH::PrintLine(FileText);

    //3. compile the shader
    auto CompileResult=RApplication::GetApplication()
        ->GetSubsystem<RShaderCompileSubsystem>()
        ->Compile(FileText,EShaderStage::Fragment);
    if(CompileResult.ErrorMessage.HasValue())
    {
        WH::PrintLine(U"着色器编译错误" + CompileResult.ErrorMessage.GetValue());
        return;
    }
    
    auto Shader=RShader::CreateShader(std::move(CompileResult.ByteCode),EShaderStage::Fragment);
    Shader->SetOwner(InMaterial.Get());

    //4.update the material
    InMaterial->Type=CurrentOutputNode->GetMaterialType();
    InMaterial->FragmentShader=Shader;
    InMaterial->Texture2Ds=Texture2Ds;
    InMaterial->FloatParameters=FloatParameters;
    InMaterial->ParameterDescriptions=ParameterDescriptions;
    InMaterial->SyncToRenderThread();

}

bool RMaterialGraphModel::HasUnsavedData() const
{
    if(LastOperationIndex<0
        && !OperationBeforeLastSave)
    {
        return false;
    }

    if(LastOperationIndex>=0
        && Operations[LastOperationIndex]==OperationBeforeLastSave)
    {
        return false;
    }

    return true;
}

void RMaterialGraphModel::Save(TSharedObjectPtr<RMaterial> InMaterial)
{
    if(LastOperationIndex<0)
    {
        OperationBeforeLastSave=nullptr;
    }
    else
    {
        OperationBeforeLastSave=Operations[LastOperationIndex];
    }

    //compile the shader
    RegenerateMaterial(InMaterial);

    //save the material
    RAssetSubsystem::Get()->ResaveLoadedAsset(
        InMaterial
        ,this
        ,nullptr);
}

void RMaterialGraphModel::SetSelectedNode(TSharedObjectPtr<RMaterialGraphNodeModel> Node)
{
    if(SelectedNode==Node)
    {
        return;
    }

    //选择节点仅仅是告诉用户当前可以删除的节点
    //不应当参与到撤销操作中
    SelectedNode=Node;
    OnSelectedNodeChanged.ExecuteIfBound( SelectedNode);


    //如果选中的是输出节点，那么将其设置为当前输出节点
    if(SelectedNode&&SelectedNode->IsOutputNode())
    {
        SetCurrentOutputNode(SelectedNode);
    }
}

void RMaterialGraphModel::RemoveSelectedNode()
{
    if(SelectedNode)
    {
        RemoveNode(SelectedNode);
    }
}

TSharedObjectPtr<RMaterialGraphNodeModel> RMaterialGraphModel::GetSelectedNode() const
{
    return SelectedNode;
}

void RMaterialGraphModel::SetCurrentOutputNode(TSharedObjectPtr<RMaterialGraphNodeModel> Node)
{
    if(CurrentOutputNode==Node)
    {
        return;
    }

    if(!CurrentOutputNode.IsValid()
        || !Node.IsValid()
    )
    {
        //由于必须有一个输出节点
        //当当前输出节点为空时 或者 新的输出节点为空时
        //不应当参与到撤销操作中
        LowLevelSetOutputNode(Node);
        return;
    }

    auto Operation=NewObject<RMaterialSetOutputNodeOperation>();
    Operation->OldOutputNode=CurrentOutputNode;
    Operation->NewOutputNode=Node;

    AddOperation(Operation);
    Operation->Redo(this);
}

TSharedObjectPtr<RMaterialGraphNodeModel> RMaterialGraphModel::GetCurrentOutputNode() const
{
    return CurrentOutputNode;
}

void RMaterialGraphModel::AddOperation(TSharedObjectPtr<RMaterialEditOperation> Operation)
{
    if(LastOperationIndex<Operations.Num()-1)
    {
        Operations.RemoveAfter(LastOperationIndex,false);
    }

    if(!Operations.Empty()  && Operation->CombineTo(Operations.Last()))
    {
        //do nothing
    }
    else
    {
        Operations.Add(Operation);
    }
    
    LastOperationIndex=Operations.Num()-1;

}

void RMaterialGraphModel::LowLevelInsertNode(TSharedObjectPtr<RMaterialGraphNodeModel> Node, int64_t Index)
{
    Nodes.Insert(Index,Node);
    OnNodeAdded.ExecuteIfBound(Index);
}

void RMaterialGraphModel::LowLevelRemoveNode(int64_t Index)
{
    auto RemovedNode=Nodes[Index];
    if(RemovedNode==SelectedNode)
    {
        SelectedNode=nullptr;
    }

    Nodes.RemoveAt(Index);
    OnNodeRemoved.ExecuteIfBound(Index);
}

void RMaterialGraphModel::LowLevelInsertLink(TSharedObjectPtr<RMaterialGraphLinkModel> Link, int64_t Index)
{
    Links.Insert(Index,Link);
    
    //没有指定连接的索引
    //可能会导致连接的顺序变化
    Link->Start->AddLink(Link);
    Link->End->AddLink(Link);

    OnLinkAdded.ExecuteIfBound(Index);
}

void RMaterialGraphModel::LowLevelRemoveLink(int64_t Index)
{
    auto RemovedLink=Links[Index];
    Links.RemoveAt(Index);

    //没有指定连接的索引
    //可能会导致连接的顺序变化
    RemovedLink->Start->RemoveLink(RemovedLink);
    RemovedLink->End->RemoveLink(RemovedLink);

    OnLinkRemoved.ExecuteIfBound(Index);
}

void RMaterialGraphModel::LowLevelSetOutputNode(TSharedObjectPtr<RMaterialGraphNodeModel> Node)
{
    CurrentOutputNode=Node;
    OnCurrentOutputNodeChanged.ExecuteIfBound(Node);
}

void RMaterialAddNodeOperation::Undo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelRemoveNode(NodeIndex);
}

void RMaterialAddNodeOperation::Redo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    Node->Position=Position;
    InGraph->LowLevelInsertNode(Node,NodeIndex);
}

void RMaterialRemoveNodeOperation::Undo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelInsertNode(Node,NodeIndex);
}

void RMaterialRemoveNodeOperation::Redo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelRemoveNode(NodeIndex);
}

void RMaterialAddLinkOperation::Undo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelRemoveLink(LinkIndex);
}

void RMaterialAddLinkOperation::Redo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelInsertLink(AddedLink,LinkIndex);
}

void RMaterialRemoveLinkOperation::Undo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelInsertLink(RemovedLink,LinkIndex);
}

void RMaterialRemoveLinkOperation::Redo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelRemoveLink(LinkIndex);
}

void RMaterialSetOutputNodeOperation::Undo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelSetOutputNode(OldOutputNode);
}

void RMaterialSetOutputNodeOperation::Redo(TSharedObjectPtr<RMaterialGraphModel> InGraph)
{
    InGraph->LowLevelSetOutputNode(NewOutputNode);
}
